# -*- coding: utf-8 -*-
import os
import numpy as np
import matplotlib.pyplot as plt
import wrf_output
import gauge
import roc_calculator


def decumulation(seq):
    """de-accumulation"""
    result = [seq[i + 1] - seq[i] for i in range(len(seq) - 1)]
    result.insert(0, seq[0])
    return result


def plot_roc(dat_path, domain, start_time, end_time, event):
    """plot roc curves"""
    # domains = ['01', '02']
    # get every mp_cu group wrf_output and extract all 50 gauge-located value
    mps = ['2', '6', '7']
    cus = ['0', '1', '2', '7']
    pbls = ['2', '6', '9']
    split_path = dat_path.split('/')
    gridSpacing = split_path[-1]

    gauge_loc = ''
    if gridSpacing == '139':
        gauge_loc = r'F:/research/rainfall_estimation/dat/Gauge' \
                    r'/Gauge_row_col_1.csv'
    elif gridSpacing == '51545':
        gauge_loc = r'F:/research/rainfall_estimation/dat/Gauge' \
                    r'/Gauge_row_col_5.csv'
    elif gridSpacing == '103090':
        gauge_loc = r'F:/research/rainfall_estimation/dat/Gauge' \
                    r'/Gauge_row_col_10.csv'

    simulations = []
    for i in mps:
        for j in cus:
            for k in pbls:
                mp_cu_pbl = i + j + k
                if os.path.exists(dat_path + '/' + mp_cu_pbl):
                    input_path = dat_path + '/' + mp_cu_pbl + \
                        '/wrfout_d' + domain + '_' + start_time
                    all_gauges = wrf_output.GetAll(input_path, gauge_loc)
                    all_gauges = decumulation(all_gauges)
                    simulations.append(all_gauges)
    simulations = np.asarray(simulations)

    # get the center gauge value
    gauge50 = gauge.Get50Gauge(start_time, end_time)

    # set threshold and probabilities
    # gauge_maximum = np.max(gauge50, axis=0)
    # max_threshold = np.mean(gauge_maximum)
    # if max_threshold > 2:
    #     thresholds = np.linspace(1, max_threshold / 2, 5)
    # else:
    #     thresholds = np.linspace(0, max_threshold / 2, 5)
    thresholds = (1.0, 1.25, 1.5, 1.75, 2.0)
    probabilities = (0.9, 0.7, 0.5, 0.3, 0.1)
    pods = []
    fars = []
    for threshold in thresholds:
        pod_thd, far_thd = [], []
        for probability in probabilities:
            # calculate pod, far, bias
            # roc_calculator.test(
            #     gauge50.values, results, threshold, probability)
            # sum up all 50 a, c, b, d to calculate pod far
            pod_prob = roc_calculator.calculate_pod(
                gauge50.values, simulations, threshold, probability)
            far_prob = roc_calculator.calculate_far(
                gauge50.values, simulations, threshold, probability)
            pod_thd.append(pod_prob)
            far_thd.append(far_prob)
        pods.append(pod_thd)
        fars.append(far_thd)
        # bias.append(bia_thd)
    # plot roc curves of center
    # labels = [str(thd) for thd in thresholds]
    labels = list(map(str, thresholds))

    # use center gauge to calculate roc
    # pods = np.array(pods)
    # fars = np.array(fars)
    # pod = np.nanmean(pods, axis=2)
    # far = np.nanmean(fars, axis=2)
    # pod = pods[:, :, 4]
    # far = fars[:, :, 4]
    # roc = np.array([far, pod])

    # use all gauges to calculate roc
    roc = np.array([fars, pods])

    start_points = np.zeros((2, 5))
    roc = np.insert(roc, 0, values=start_points, axis=2)
    end_points = np.ones((2, 5, 1))
    roc = np.c_[roc, end_points]
    thd_num = roc.shape[1]

    plt.style.use('fivethirtyeight')
    fig, ax = plt.subplots(figsize=(10, 10))
    color_seq = ['#1f77b4', '#2ca02c', '#bcbd22', '#ff7f0e', '#d62728']
    for j in range(thd_num):
        far_pod = roc[:, j, :]
        ax.plot(far_pod[0], far_pod[1], color=color_seq[j], label=labels[j])
        roc_area = roc_calculator.calculate_roc_area(far_pod[0], far_pod[1])
        print(roc_area)
    ax.plot([0, 1], [0, 1], '--', color='gray', label='no-skill')
    # plt.legend(loc='lower right', fontsize='xx-large')

    # adjust label
    plt.xlabel('FAR', fontdict={'size': 22})
    # plt.ylabel('POD', fontdict={'size': 22})
    ax.set_yticklabels([])

    for x_label in plt.gca().xaxis.get_ticklabels():
        x_label.set_fontsize(20)
    for y_label in plt.gca().yaxis.get_ticklabels():
        y_label.set_fontsize(20)
    plt.title(event, fontsize=24)
    plt.tight_layout()
    plt.show()
    # gauge_num = pods.shape[2]
    # for i in range(gauge_num):
    #     pod = pods[:, :, i]
    #     far = fars[:, :, i]
    #     roc = np.array([far, pod])
    #     start_points = np.zeros((2, 5))
    #     end_points = np.ones((2, 5, 1))
    #     roc = np.insert(roc, 0, values=start_points, axis=2)
    #     roc = np.c_[roc, end_points]
    #     thd_num = roc.shape[1]
    #     plt.figure(figsize=(10, 10))
    #     for j in range(thd_num):
    #         far_pod = roc[:, j, :]
    #         far_pod_arg = np.argsort(far_pod[0, :])
    #         far_pod_sorted = far_pod[:, far_pod_arg]
    #         plt.plot(far_pod_sorted[0], far_pod_sorted[1], label=labels[j])
    #     plt.plot([0, 1], [0, 1], 'r--', label='no-skill')
    #     plt.legend(loc='best', fontsize='xx-large')
    # plt.show()


def main():
    """main function"""
    # new run
    # R1
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/01/1700/103090'
    # domain = r'03'
    # start_time = r'2008-01-17_00_00_00'
    # end_time = r'2008-01-19_12_00_00'
    # event = 'R1'

    # R2
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/01/1912/103090'
    # domain = r'03'
    # start_time = r'2008-01-19_12_00_00'
    # end_time = r'2008-01-22_00_00_00'
    # event = 'R2'

    # R3
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/08/1718/103090'
    # domain = r'03'
    # start_time = r'2008-08-17_18_00_00'
    # end_time = r'2008-08-20_00_00_00'
    # event = 'R3'

    # R4
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/09/0500/103090'
    # domain = r'03'
    # start_time = r'2008-09-05_00_00_00'
    # end_time = r'2008-09-07_00_00_00'
    # event = 'R4'

    # R5
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/09/2900/103090'
    # domain = r'03'
    # start_time = r'2008-09-29_00_00_00'
    # end_time = r'2008-10-02_06_00_00'
    # event = 'R5'

    # R6
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/10/2506/103090'
    # domain = r'03'
    # start_time = r'2008-10-25_06_00_00'
    # end_time = r'2008-10-26_06_00_00'
    # event = 'R6'

    # R7
    # dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/11/0900/103090'
    # domain = r'03'
    # start_time = r'2008-11-09_00_00_00'
    # end_time = r'2008-11-10_06_00_00'
    # event = 'R7'

    # R8
    dat_path = r'H:/research/rainfall_estimation/wrf_output/2008/12/0400/103090'
    domain = r'03'
    start_time = r'2008-12-04_00_00_00'
    end_time = r'2008-12-06_00_00_00'
    event = 'R8'

    plot_roc(dat_path, domain, start_time, end_time, event)


if __name__ == '__main__':
    main()
